{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "from ipywidgets import interact\n",
    "import matplotlib.pyplot as plt\n",
    "from chempy import Reaction, Substance, ReactionSystem\n",
    "from chempy.kinetics.ode import get_odesys\n",
    "from chempy.kinetics.analysis import plot_reaction_contributions\n",
    "from chempy.printing.tables import UnimolecularTable, BimolecularTable\n",
    "from chempy.util.graph import rsys2graph\n",
    "import sympy\n",
    "sympy.init_printing()\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A, B, C, D = map(Substance, 'ABCD')\n",
    "One = sympy.S.One\n",
    "reactions = r0, r1, r2 = [\n",
    "    Reaction({'A'}, {'B'}, 4*One/100, name='R1: A cons.'),\n",
    "    Reaction({'B', 'C'}, {'A', 'C'}, 10**(4*One), name='R2: A reform.'),\n",
    "    Reaction({'B': 2}, {'B', 'C'}, 3*10**(7*One), name='R3: C form.')\n",
    "]\n",
    "rsys = ReactionSystem(reactions, (A, B, C))\n",
    "rsys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uni, not_uni = UnimolecularTable.from_ReactionSystem(rsys)\n",
    "bi, not_bi = BimolecularTable.from_ReactionSystem(rsys)\n",
    "assert not (not_uni & not_bi), \"There are only uni- & bi-molecular reactions in this set\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uni"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rsys2graph(rsys, 'robertson.png', save='.')\n",
    "from IPython.display import Image; Image('robertson.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "odesys, extra = get_odesys(rsys, include_params=True)\n",
    "odesys.exprs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "odesys.get_jac()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c0 = defaultdict(float, {'A': 1})\n",
    "result = odesys.integrate(1e10, c0, integrator='cvode', nsteps=2000)\n",
    "{k: v for k, v in result.info.items() if not k.startswith('internal')}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "extra['rate_exprs_cb'](result.xout, result.yout)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.plot(xscale='log', yscale='log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2, 2, figsize=(14, 6))\n",
    "plot_reaction_contributions(result, rsys, extra['rate_exprs_cb'], 'AB', axes=axes[0, :])\n",
    "plot_reaction_contributions(result, rsys, extra['rate_exprs_cb'], 'AB', axes=axes[1, :],\n",
    "                            relative=True, yscale='linear')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could also have parsed the reactions from a string:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "str_massaction = \"\"\"\n",
    "A -> B; 'k1'\n",
    "B + C -> A + C; 'k2'\n",
    "2 B -> B + C; 'k3'\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rsys3 = ReactionSystem.from_string(str_massaction, substance_factory=lambda formula: Substance(formula))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rsys3.substance_names()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "odesys3, extra3 = get_odesys(rsys3, include_params=False, lower_bounds=[0, 0, 0])\n",
    "extra3['param_keys'], extra3['unique']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "odesys3.exprs, odesys3.params, odesys3.names, odesys3.param_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def integrate_and_plot(A0=1.0, B0=0.0, C0=0.0, lg_k1=-2, lg_k2=4, lg_k3=7, lg_tend=9):\n",
    "    plt.figure(figsize=(14, 4))\n",
    "    tout, yout, info = odesys3.integrate(\n",
    "        10**lg_tend, {'A': A0, 'B': B0, 'C': C0},\n",
    "        {'k1': 10**lg_k1, 'k2': 10**lg_k2, 'k3': 10**lg_k3},\n",
    "        integrator='cvode', nsteps=3000)\n",
    "    plt.subplot(1, 2, 1)\n",
    "    odesys3.plot_result(xscale='log', yscale='log')\n",
    "    plt.legend(loc='best')\n",
    "    plt.subplot(1, 2, 2)\n",
    "    plt.plot(tout[tout<.05], yout[tout<.05, odesys3.names.index('B')])\n",
    "    _ = plt.legend('best')\n",
    "interact(integrate_and_plot) #, **kw)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We could also have used SymPy to construct symbolic rates:\n",
    "import sympy\n",
    "rsys_sym = ReactionSystem.from_string(\"\"\"\n",
    "A -> B; sp.Symbol('k1')\n",
    "B + C -> A + C; sp.Symbol('k2')\n",
    "2 B -> B + C; sp.Symbol('k3')\n",
    "\"\"\", rxn_parse_kwargs=dict(globals_={'sp': sympy}), substance_factory=lambda formula: Substance(formula))\n",
    "odesys_sym, _ = get_odesys(rsys_sym, params=True)\n",
    "for attr in 'exprs params names param_names'.split():\n",
    "    print(getattr(odesys_sym, attr))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For larger systems it is easy to loose track of what substances are actually playing a part, here the html tables can help (note the yellow background color):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rsys.substances['D'] = D\n",
    "uni, not_uni = UnimolecularTable.from_ReactionSystem(rsys)\n",
    "uni"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bi, not_bi = BimolecularTable.from_ReactionSystem(rsys)\n",
    "bi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0+"
  },
  "widgets": {
   "state": {
    "6296ca2fc06a4f0c9d9390c4dd51a493": {
     "views": [
      {
       "cell_index": 16
      }
     ]
    }
   },
   "version": "1.2.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
